// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#include <pollux/serializers/presto_serializer.h>

#include <optional>

#include <melon/lang/bits.h>

#include <pollux/common/base/crc.h>
#include <pollux/common/memory/byte_stream.h>
#include <pollux/serializers/presto_batch_vector_serializer.h>
#include <pollux/serializers/presto_header.h>
#include <pollux/serializers/presto_iterative_vector_serializer.h>
#include <pollux/serializers/presto_serializer_deserialization_utils.h>
#include <pollux/serializers/presto_serializer_estimation_utils.h>
#include <pollux/serializers/presto_serializer_serialization_utils.h>
#include <pollux/serializers/presto_vector_lexer.h>
#include <pollux/serializers/vector_stream.h>
#include <pollux/vector/complex_vector.h>

namespace kumo::pollux::serializer::presto {
    using SerdeOpts = PrestoVectorSerde::PrestoOptions;

    namespace {
        int64_t computeChecksum(
            ByteInputStream *source,
            int codecMarker,
            int numRows,
            int32_t uncompressedSize,
            int32_t compressedSize) {
            const auto offset = source->tellp();
            const bool compressed = codecMarker & detail::kCompressedBitMask;
            if (compressed) {
                POLLUX_CHECK_LT(compressedSize, uncompressedSize);
            }
            const int32_t dataSize = compressed ? compressedSize : uncompressedSize;
            bits::Crc32 crc32;
            if (MELON_UNLIKELY(source->remainingSize() < dataSize)) {
                POLLUX_FAIL(
                    "Tried to read {} bytes, larger than what's remained in source {} "
                    "bytes. Source details: {}",
                    dataSize,
                    source->remainingSize(),
                    source->toString());
            }
            auto remainingBytes = dataSize;
            while (remainingBytes > 0) {
                auto data = source->nextView(remainingBytes);
                if (MELON_UNLIKELY(data.size() == 0)) {
                    POLLUX_FAIL(
                        "Reading 0 bytes from source. Source details: {}",
                        source->toString());
                }
                crc32.process_bytes(data.data(), data.size());
                remainingBytes -= data.size();
            }

            crc32.process_bytes(&codecMarker, 1);
            crc32.process_bytes(&numRows, 4);
            crc32.process_bytes(&uncompressedSize, 4);
            auto checksum = crc32.checksum();

            source->seekp(offset);

            return checksum;
        }

        PrestoVectorSerde::PrestoOptions toPrestoOptions(
            const VectorSerde::Options *options) {
            if (options == nullptr) {
                return PrestoVectorSerde::PrestoOptions();
            }
            auto prestoOptions =
                    dynamic_cast<const PrestoVectorSerde::PrestoOptions *>(options);
            POLLUX_CHECK_NOT_NULL(
                prestoOptions, "Serde options are not Presto-compatible");
            return *prestoOptions;
        }
    } // namespace

    void PrestoVectorSerde::estimateSerializedSize(
        const BaseVector *vector,
        const melon::Range<const IndexRange *> &ranges,
        vector_size_t **sizes,
        Scratch &scratch) {
        detail::estimateSerializedSizeInt(
            vector->loaded_vector(), ranges, sizes, scratch);
    }

    void PrestoVectorSerde::estimateSerializedSize(
        const BaseVector *vector,
        const melon::Range<const vector_size_t *> &rows,
        vector_size_t **sizes,
        Scratch &scratch) {
        detail::estimateSerializedSizeInt(
            vector->loaded_vector(), rows, sizes, scratch);
    }

    std::unique_ptr<IterativeVectorSerializer>
    PrestoVectorSerde::createIterativeSerializer(
        RowTypePtr type,
        int32_t numRows,
        StreamArena *streamArena,
        const Options *options) {
        const auto prestoOptions = toPrestoOptions(options);
        return std::make_unique<detail::PrestoIterativeVectorSerializer>(
            type, numRows, streamArena, prestoOptions);
    }

    std::unique_ptr<BatchVectorSerializer> PrestoVectorSerde::createBatchSerializer(
        memory::MemoryPool *pool,
        const Options *options) {
        const auto prestoOptions = toPrestoOptions(options);
        return std::make_unique<detail::PrestoBatchVectorSerializer>(
            pool, prestoOptions);
    }

    void PrestoVectorSerde::deserialize(
        ByteInputStream *source,
        pollux::memory::MemoryPool *pool,
        RowTypePtr type,
        RowVectorPtr *result,
        vector_size_t resultOffset,
        const Options *options) {
        const auto prestoOptions = toPrestoOptions(options);
        const auto codec =
                common::compressionKindToCodec(prestoOptions.compressionKind);
        auto maybeHeader = detail::PrestoHeader::read(source);
        POLLUX_CHECK(
            maybeHeader.hasValue(),
            fmt::format(
                "PrestoPage header is invalid: {}", maybeHeader.error().message()));
        auto const header = std::move(maybeHeader.value());

        int64_t actualCheckSum = 0;
        if (detail::isChecksumBitSet(header.pageCodecMarker)) {
            actualCheckSum = computeChecksum(
                source,
                header.pageCodecMarker,
                header.numRows,
                header.uncompressedSize,
                header.compressedSize);
        }

        POLLUX_CHECK_EQ(
            header.checksum, actualCheckSum, "Received corrupted serialized page.");

        if (resultOffset > 0) {
            POLLUX_CHECK_NOT_NULL(*result);
            POLLUX_CHECK_EQ(result->use_count(), 1);
            (*result)->resize(resultOffset + header.numRows);
        } else if (*result && result->use_count() == 1) {
            POLLUX_CHECK(
                *(*result)->type() == *type,
                "Unexpected type: {} vs. {}",
                (*result)->type()->toString(),
                type->toString());
            (*result)->prepare_for_reuse();
            (*result)->resize(header.numRows);
        } else {
            *result = BaseVector::create<RowVector>(type, header.numRows, pool);
        }

        POLLUX_CHECK_EQ(
            header.checksum, actualCheckSum, "Received corrupted serialized page.");

        if (!detail::isCompressedBitSet(header.pageCodecMarker)) {
            detail::readTopColumns(
                *source, type, pool, *result, resultOffset, prestoOptions);
        } else {
            auto compressBuf = melon::IOBuf::create(header.compressedSize);
            source->readBytes(compressBuf->writableData(), header.compressedSize);
            compressBuf->append(header.compressedSize);

            // Process chained uncompressed results IOBufs.
            auto uncompress =
                    codec->uncompress(compressBuf.get(), header.uncompressedSize);
            auto uncompressedSource = std::make_unique<BufferInputStream>(
                byteRangesFromIOBuf(uncompress.get()));
            detail::readTopColumns(
                *uncompressedSource, type, pool, *result, resultOffset, prestoOptions);
        }
    }

    void PrestoVectorSerde::deserializeSingleColumn(
        ByteInputStream *source,
        pollux::memory::MemoryPool *pool,
        TypePtr type,
        VectorPtr *result,
        const Options *options) {
        const auto prestoOptions = toPrestoOptions(options);
        POLLUX_CHECK_EQ(
            prestoOptions.compressionKind,
            common::CompressionKind::CompressionKind_NONE);
        if (*result && result->use_count() == 1) {
            POLLUX_CHECK(
                *(*result)->type() == *type,
                "Unexpected type: {} vs. {}",
                (*result)->type()->toString(),
                type->toString());
            (*result)->prepare_for_reuse();
        } else {
            *result = BaseVector::create(type, 0, pool);
        }

        auto rowType = ROW({"c0"}, {type});
        auto row = std::make_shared<RowVector>(
            pool, rowType, BufferPtr(nullptr), 0, std::vector<VectorPtr>{*result});
        detail::readTopColumns(*source, rowType, pool, row, 0, prestoOptions, true);
        *result = row->childAt(0);
    }

    void PrestoVectorSerde::serializeSingleColumn(
        const VectorPtr &vector,
        const Options *opts,
        memory::MemoryPool *pool,
        std::ostream *output) {
        const auto prestoOptions = toPrestoOptions(opts);
        POLLUX_USER_CHECK_EQ(
            prestoOptions.compressionKind,
            common::CompressionKind::CompressionKind_NONE);
        POLLUX_USER_CHECK_EQ(prestoOptions.nullsFirst, false);

        const IndexRange range{0, vector->size()};
        const auto arena = std::make_unique<StreamArena>(pool);
        auto stream = std::make_unique<detail::VectorStream>(
            vector->type(),
            std::nullopt,
            std::nullopt,
            arena.get(),
            vector->size(),
            prestoOptions);
        Scratch scratch;
        serializeColumn(vector, melon::Range(&range, 1), stream.get(), scratch);

        PrestoOutputStreamListener listener;
        OStreamOutputStream outputStream(output, &listener);
        stream->flush(&outputStream);
    }

    // static
    void PrestoVectorSerde::registerVectorSerde() {
        detail::initBitsToMapOnce();
        pollux::registerVectorSerde(std::make_unique<PrestoVectorSerde>());
    }

    // static
    void PrestoVectorSerde::registerNamedVectorSerde() {
        detail::initBitsToMapOnce();
        pollux::registerNamedVectorSerde(
            VectorSerde::Kind::kPresto, std::make_unique<PrestoVectorSerde>());
    }

    /* static */
    Status PrestoVectorSerde::lex(
        std::string_view source,
        std::vector<Token> &out,
        const Options *options) {
        const auto prestoOptions = toPrestoOptions(options);

        POLLUX_RETURN_IF(
            prestoOptions.useLosslessTimestamp,
            Status::Invalid(
                "Lossless timestamps are not supported, because they cannot be decoded without the Schema"));
        POLLUX_RETURN_IF(
            prestoOptions.compressionKind !=
            common::CompressionKind::CompressionKind_NONE,
            Status::Invalid("Compression is not supported"));
        POLLUX_RETURN_IF(
            prestoOptions.nullsFirst,
            Status::Invalid(
                "Nulls first encoding is not currently supported, but support can be added if needed"));

        return detail::PrestoVectorLexer(source).lex(out);
    }
} // namespace kumo::pollux::serializer::presto
